(library (grid-gen math)
  (export math-rvector
          math-cvector
          math-array
          math-array-el
          math-one
          arr-elements
          arr-row
          arr-col
          arr-cols
          diag
          add transp mul2 abs2 arr-max
          math-array-dims
          arr-n->ind
          arr-ind->n
          map-ar
          infix-out
          swap-col
          swap-str
          levi-civita
          all-inds
  )

;;(use-modules (statprof))

(import ;;(rnrs)
        ;;(rnrs r5rs)
        (grid-gen main)
        (grid-gen vars)
        ;;(grid-gen symath)
        ;;(chezscheme)
        (guile)
        (except (srfi srfi-1) map)
)

;;(import (srfi srfi-1))
;;(import (ice-9 threads))

(define math-array (lambda (dims itms)
  (list 'M_ARR (list 'M_DIM dims) itms)
))

(define math-one (lambda (s)
  (letrec* ((nmax (* s s))
            (dims (list s s))
            (a-i (lambda (n)
                   (let ((ind (arr-n->ind n dims)))
                     (if (= n nmax)
                         (list)
                         (cons (if (= (car ind) (cadr ind))
                                   1.0
                                   0.0
                               )
                               (a-i (+ 1 n))
                         )
                     )
                   )
                 )
            )
           )
    (math-array dims (a-i 0))
  )
))

(define diag (lambda (v)
  (letrec* ((s (cadr (math-array-dims v)))
            (els (arr-elements v))
            (nmax (* s s))
            (dims (list s s))
            (a-i (lambda (n)
                   (let ((ind (arr-n->ind n dims)))
                     (if (= n nmax)
                         (list)
                         (cons (if (= (car ind) (cadr ind))
                                   (list-ref els (car ind))
                                   0.0
                               )
                               (a-i (+ 1 n))
                         )
                     )
                   )
                 )
            )
           )
    (math-array dims (a-i 0))
  )
))

(define is-math-array? (lambda (v)
  (and (list? v) (not (null? v)) 
       (eq? (car v) 'M_ARR)
  )
))

(define math-array-dims (lambda (a)
  (letrec ((m-a-d (lambda (l)
                    (if (not (null? l))
                        (if (and (list? (car l))
                                 (not (null? (car l)))
                                 (eq? 'M_DIM (caar l))
                            )
                            (cadar l)
                            (m-a-d (cdr l))
                        )
                        (int-err "Math array dims not found")
                    )
                  )
          ))
    (if (is-math-array? a)
        (m-a-d (cdr a))
        (int-err "Not an array")
    )
  )
))

(define arr-elements (lambda (a)
  (list-ref a 2)
))

(define math-rvector (lambda itms (math-array (list (length itms) 1) itms)))
(define math-cvector (lambda itms (math-array (list 1 (length itms)) itms)))

(define arr-add (lambda (l1 l2)
  (if (and (null? l1) (null? l2))
      (list)
      (cons (cond ((and (number? (car l1)) (= (car l1) 0.0)) (car l2))
                  ((and (number? (car l2)) (= (car l2) 0.0)) (car l1))
                  ((and (number? (car l1)) (number? (car l2))) 
                   (+ (car l1) (car l2))
                  )
                  (#t `(+ ,(car l1) ,(car l2)))
            )
            (arr-add (cdr l1) (cdr l2))
      )
  )
))

(define is-scalar? (lambda (v)
  (not (is-math-array? v))
))

(define add (lambda args
  (if (null? (cdr args))
      (car args)
      (let ((rst (apply add (cdr args))))
        (cond ((and (is-math-array? (car args))
                    (is-math-array? rst)
                    (equal? (math-array-dims (car args))
                            (math-array-dims rst)
                    )
               )
               (math-array (math-array-dims rst)
                           (arr-add (arr-elements (car args)) 
                                    (arr-elements rst)
                           )
               )
              )
              ((and (is-scalar? (car args)) (is-scalar? rst))
               (cond
                 ((and (number? (car args)) (number? rst)) (+ (car args) rst))
                 ((and (number? (car args)) (= (car args) 0)) rst)
                 ((and (number? rst) (= rst 0)) (car args))
                 (#t `(+ ,(car args) ,rst))
               )
              )
              (#t (begin
                   ;; (format #t "A: ~a\n" args)
                    (int-err "Unsupported types for add")
                  )
              )
        )
      )
  )
))

(define arr-ind->n (lambda (i d)
  (if (null? i)
    0
    (+ (car i) (* (car d) (arr-ind->n (cdr i) (cdr d))))
  )
))

(define arr-n->ind (lambda (n d)
  (if (null? d)
      (list)
      (let* ((m (car d))
             (i (modulo n m))
             (j (floor (/ n m)))
            )
        (cons i (arr-n->ind j (cdr d)))
      )
  )
))

(define swap-str (lambda (v c1 c2)
  (if (is-math-array? v)
      (letrec* ((els (arr-elements v))
                (dims (math-array-dims v))
                (l-p (lambda (n)
                  (if (= n (length els))
                      (list)
                      (cons
                        (let* ((i (arr-n->ind n dims))
                               (i1 (list (car i)
                                         (cond ((= (cadr i) c1) c2)
                                               ((= (cadr i) c2) c1)
                                               (#t (cadr i))
                                         )
                                   )
                               )
                              )
                          (list-ref els (arr-ind->n i1 dims))
                        )
                        (l-p (+ n 1))
                      )
                  )
                ))
               )
        (math-array dims (l-p 0))
      )
      (int-err "Not an array in swap-col")
  )
))

(define swap-col (lambda (v s1 s2)
  (if (is-math-array? v)
      (letrec* ((els (arr-elements v))
                (dims (math-array-dims v))
                (l-p (lambda (n)
                  (if (= n (length els))
                      (list)
                      (cons
                        (let* ((i (arr-n->ind n dims))
                               (i1 (list (cond ((= (car i) s1) s2)
                                               ((= (car i) s2) s1)
                                               (#t (car i))
                                         )
                                         (cadr i)
                                   )
                               )
                              )
                          (list-ref els (arr-ind->n i1 dims))
                        )
                        (l-p (+ n 1))
                      )
                  )
                ))
               )
        (math-array dims (l-p 0))
      )
      (int-err "Not an array in swap-str")
  )
))

(define transp (lambda (v)
  (if (is-math-array? v)
      (letrec ((dims (math-array-dims v))
               (els  (arr-elements v))
               (arr-t (lambda (n)
                        (if (= n (length els))
                            (list)
                            (let ((i (arr-n->ind n dims)))
                              (cons (list-ref els
                                              (arr-ind->n (roll i)
                                                          (roll dims)
                                              )
                                    )
                                    (arr-t (+ 1 n))
                              )
                            )
                        )
                      )
               )
              )
        (math-array (roll dims) (arr-t 0))
      )
      v
  )
))

(define mul-arr-scal (lambda (a s)
  (cond ((and (number? s) (= 0.0 s))
         (math-array (math-array-dims a)
                     (map (lambda (x) 0.0) (arr-elements a))
         )
        )
        ((and (number? s) (= 1.0 s)) a)
        (#t (math-array
              (math-array-dims a)
              (map (lambda (x) 
                     (cond ((and (number? x) (= 0.0 x)) 0.0)
                           ((and (number? x) (= 1.0 x)) s)
                           ((and (number? x) (number? s)) (* x s))
                           (#t `(* ,s ,x))
                     )
                   )
                   (arr-elements a)
              )
            )
        )
  )
))

(define mul-arr-arr (lambda (a1 a2)
  (letrec* ((d1 (math-array-dims a1))
            (d2 (math-array-dims a2))
            (rdims (list (car d2) (cadr d1))) ;;;; 2D ONLY ;;;;
            (e1 (arr-elements a1))
            (e2 (arr-elements a2))
            (f-a (lambda (i)
              (if (= i (apply * rdims))
                (list)
                (cons 
                  (letrec* ((ind (arr-n->ind i rdims))
                            (f-a-l (lambda (j)
                              (if (= j (cadr rdims))
                                0.0
                                (let* ((t1 (list-ref 
                                             e1
                                             (arr-ind->n (list j (cadr ind))
                                                         d1
                                             )
                                           )
                                       )
                                       (t2 (list-ref 
                                             e2
                                             (arr-ind->n (list (car ind) j)
                                                         d2
                                             )
                                           )
                                       )
                                       (t3 (f-a-l (+ j 1)))
                                       (t12 (cond ((and (number? t1)
                                                        (= t1 0.0)
                                                   )
                                                   0.0
                                                  )
                                                  ((and (number? t2)
                                                        (= t2 0.0)
                                                   )
                                                   0.0
                                                  )
                                                  ((and (number? t1)
                                                        (= t1 1.0)
                                                   )
                                                   t2
                                                  )
                                                  ((and (number? t2)
                                                        (= t2 1.0)
                                                   )
                                                   t1
                                                  )
                                                  ((and (number? t1)
                                                        (number? t2)
                                                   )
                                                   (* t1 t2)
                                                  )
                                                  (#t `(* ,t1 ,t2))
                                            )
                                       )
                                      )
                                  (cond ((and (number? t12) (= t12 0.0)) t3)
                                        ((and (number? t3) (= t3 0.0)) t12)
                                        ((and (number? t12) (number? t3))
                                         (+ t12 t3)
                                        )
                                        (#t `(+ ,t12 ,t3))
                                  )
                                )
                              )
                            ))
                           )
                    (f-a-l 0)
                  )
                  (f-a (+ i 1))
                )
              )
            ))
           )
    (math-array rdims (f-a 0))
  )
))

(define mul2 (lambda (v1 v2)
  (cond ((and (is-scalar? v1) (is-scalar? v2))
         (cond ((and (number? v1) (number? v2)) (* v1 v2))
               ((and (number? v1) (= v1 0)) 0.0)
               ((and (number? v2) (= v2 0)) 0.0)
               ((and (number? v1) (= v1 1)) v2)
               ((and (number? v2) (= v2 1)) v1)
               (#t `(* ,v1 ,v2))
         )
        )
        ((and (is-math-array? v1) (is-scalar? v2))
         (mul-arr-scal v1 v2)
        )
        ((and (is-math-array? v2) (is-scalar? v1))
         (mul-arr-scal v2 v1)
        )
        ((and (is-math-array? v1) (is-math-array? v2))
         (mul-arr-arr v1 v2)
        )
        (#t (int-err "Unsupported types for mul2"))
  )
))

(define math-array-el (lambda (ar ind)
  (list-ref (arr-elements ar) (arr-ind->n ind (math-array-dims ar)))
))

(define abs2 (lambda (v)
  (if (is-math-array? v)
      (math-array (math-array-dims v)
                  (map (lambda (x)
                         (if (number? x)
                             (abs x)
                            `(abs ,x)
                         )
                       )
                       (arr-elements v)
                  )
      )
      (if (number? v)
          (abs v)
         `(abs ,v)
      )
  )
))

(define arr-max (lambda (v)
  (if (is-math-array? v)
      (cons 'max (arr-elements v))
      (int-err "Invalid arg for arr-max")
  )
))


(define map-ar (lambda (f . args)
  (math-array (math-array-dims (car args))
              (map (lambda (al) (apply f al)) 
                   (apply zip (map arr-elements args))
              )
  )
))

(define infix-out (lambda (f)
  (cond ((symbol? f) (sym->str f))
        ((string? f) f)
        ((is-var? f) (format #f "~a_~a" (var-name f) (cadr (var-id f))))
        ((number? f) 
         (if (< f 0)
             (format #f "(~a)" (number->string f))
             (format #f "~a" (number->string f))
         )
        )
        ((and (list? f) (> (length f) 1))
         (cond ((or (eq? (car f) '*)
                    (eq? (car f) '/)
                    (eq? (car f) '+)
                    (eq? (car f) '-)
                    (eq? (car f) '>)
                    (eq? (car f) '<)
                    (eq? (car f) '>=)
                    (eq? (car f) '<=)
                    (eq? (car f) '=)
                )
                (string-append "("
                  (string-join (map infix-out (cdr f)) (sym->str (car f)))
                               ")"
                )
               )
               ((eq? (car f) 'inv)
                (format #f "(1.0/(~a))" (infix-out (cadr f)))
               )
               ((eq? (car f) '**)
                (format #f "(~a^~a)" (infix-out (cadr f))
                                     (infix-out (caddr f))
                )
               )
               ((eq? (car f) 'neg)
                (format #f "(-1.0*(~a))" (infix-out (cadr f)))
               )
               (#t (format #f "~a(~a)" 
                       (car f) 
                       (string-join (map infix-out (cdr f)) ",")
                   )
               )
         )
        )
        (#t (begin (format #t "ERR: ~a\n" f)
            (int-err "Cannot infix-out")
            )
        )
        
  )
))

(define n-perm (lambda (l1 l2)
  (if (< (length l2) 2)
      0
      (if (<= (car l2) (cadr l2))
          (n-perm (append l1 (list (car l2))) (cdr l2))
          (+ 1 (n-perm (list)
                       (append l1 (list (cadr l2) (car l2)) (cddr l2))
               )
          )
      )
  )
))

(define has-dbls (lambda (l)
  (and (not (null? l))
       (or (in-list? (car l) (cdr l))
           (has-dbls (cdr l))
       )
  )
))

(define levi-civita (lambda (i)
  (if (has-dbls i)
      0
      (if (odd? (n-perm (list) i))
         -1
          1
      )
  )
))

(define arr-row (lambda (ar r)
  (letrec*
    ((dims (math-array-dims ar))
     (els  (arr-elements ar))
     (l-p (lambda (n)
       (if (< n (apply * dims))
           (let ((i (arr-n->ind n dims)))
             (if (= (cadr i) r)
                 (cons (list-ref els n)
                       (l-p (+ n 1))
                 )
                 (l-p (+ n 1))
             )
           )
           (list)
       )
     ))
    )
    (apply math-rvector (l-p 0))
  )
))

(define arr-col (lambda (ar c)
  (letrec*
    ((dims (math-array-dims ar))
     (els  (arr-elements ar))
     (l-p (lambda (n)
       (if (< n (apply * dims))
           (let ((i (arr-n->ind n dims)))
             (if (= (car i) c)
                 (cons (list-ref els n)
                       (l-p (+ n 1))
                 )
                 (l-p (+ n 1))
             )
           )
           (list)
       )
     ))
    )
    (apply math-cvector (l-p 0))
  )
))

(define arr-cols (lambda (ar)
  (map (lambda (i) (arr-col ar i)) 
       (iota1 (car (math-array-dims ar)))
  )
))

(define all-inds (lambda (i d)
  (if (not (null? d))
      (if (< (car i) (car d))
          (append
             (all-inds (cons 0 i) (cdr d))
             (all-inds (cons (+ 1 (car i)) (cdr i)) d)
          )
          (list)
      )
      (list (cdr i))
  )
))


)

(import (grid-gen main))

;(self-cont 'M_ARR)
;(self-cont 'M_DIM)
